"""
涨停板数据获取模块

提供涨停板相关数据的获取功能，包括：
- 实时涨停板池
- 历史涨停数据
- 分时数据
- 封单强度等

作者: AI Assistant
版本: 1.0.0
日期: 2024-10-10
"""

import akshare as ak
import pandas as pd
from datetime import datetime, timedelta
import logging
from typing import Optional, List, Dict

logger = logging.getLogger(__name__)


class LimitUpDataFetcher:
    """涨停板数据获取器"""
    
    @staticmethod
    def get_today_limit_up_pool() -> Optional[pd.DataFrame]:
        """
        获取今日涨停板池
        
        返回:
            DataFrame: 涨停板数据，包含股票代码、名称、涨停时间、封板强度等
        """
        try:
            logger.info("正在获取今日涨停板池...")
            df = ak.stock_zt_pool_em(date=datetime.now().strftime("%Y%m%d"))
            
            if df is not None and not df.empty:
                # 重命名列
                column_map = {
                    '代码': 'code',
                    '名称': 'name',
                    '涨跌幅': 'change_pct',
                    '最新价': 'price',
                    '成交额': 'amount',
                    '流通市值': 'market_cap',
                    '总市值': 'total_market_cap',
                    '换手率': 'turnover_rate',
                    '封板资金': 'seal_amount',
                    '首次封板时间': 'first_limit_up_time',
                    '最后封板时间': 'last_limit_up_time',
                    '板上成交额': 'limit_up_amount',
                    '连板数': 'consecutive_limit_up',
                    '所属行业': 'industry'
                }
                
                for old_col, new_col in column_map.items():
                    if old_col in df.columns:
                        df.rename(columns={old_col: new_col}, inplace=True)
                
                logger.info(f"成功获取今日涨停板数据，共{len(df)}只股票")
                return df
            else:
                logger.warning("今日暂无涨停股票")
                return pd.DataFrame()
                
        except Exception as e:
            logger.error(f"获取今日涨停板池失败: {e}")
            return None
    
    @staticmethod
    def get_yesterday_limit_up_pool() -> Optional[pd.DataFrame]:
        """
        获取昨日涨停板池
        
        返回:
            DataFrame: 昨日涨停板数据
        """
        try:
            logger.info("正在获取昨日涨停板池...")
            df = ak.stock_zt_pool_previous_em(date=datetime.now().strftime("%Y%m%d"))
            
            if df is not None and not df.empty:
                logger.info(f"成功获取昨日涨停板数据，共{len(df)}只股票")
                return df
            else:
                logger.warning("昨日暂无涨停股票")
                return pd.DataFrame()
                
        except Exception as e:
            logger.error(f"获取昨日涨停板池失败: {e}")
            return None
    
    @staticmethod
    def get_strong_stocks() -> Optional[pd.DataFrame]:
        """
        获取强势股池
        
        返回:
            DataFrame: 强势股数据
        """
        try:
            logger.info("正在获取强势股池...")
            df = ak.stock_zt_pool_strong_em(date=datetime.now().strftime("%Y%m%d"))
            
            if df is not None and not df.empty:
                logger.info(f"成功获取强势股数据，共{len(df)}只股票")
                return df
            else:
                logger.warning("暂无强势股数据")
                return pd.DataFrame()
                
        except Exception as e:
            logger.error(f"获取强势股池失败: {e}")
            return None
    
    @staticmethod
    def get_stock_minute_data(symbol: str, period: str = "5") -> Optional[pd.DataFrame]:
        """
        获取股票分时数据
        
        参数:
            symbol: 股票代码
            period: 周期，1或5分钟
            
        返回:
            DataFrame: 分时数据
        """
        try:
            logger.info(f"正在获取{symbol}的{period}分钟数据...")
            df = ak.stock_zh_a_hist_min_em(
                symbol=symbol,
                period=period,
                adjust="qfq"
            )
            
            if df is not None and not df.empty:
                logger.info(f"成功获取{symbol}的分时数据，共{len(df)}条记录")
                return df
            else:
                logger.warning(f"{symbol}暂无分时数据")
                return pd.DataFrame()
                
        except Exception as e:
            logger.error(f"获取{symbol}分时数据失败: {e}")
            return None
    
    @staticmethod
    def calculate_seal_strength(row: pd.Series) -> float:
        """
        计算封板强度
        
        参数:
            row: 包含封单金额和流通市值的数据行
            
        返回:
            float: 封板强度百分比
        """
        try:
            if 'seal_amount' in row and 'market_cap' in row:
                seal_amount = float(row['seal_amount']) if pd.notna(row['seal_amount']) else 0
                market_cap = float(row['market_cap']) if pd.notna(row['market_cap']) else 1
                
                if market_cap > 0:
                    return (seal_amount / market_cap) * 100
            return 0.0
        except:
            return 0.0
    
    @staticmethod
    def get_limit_up_time_category(time_str: str) -> str:
        """
        获取涨停时间分类
        
        参数:
            time_str: 涨停时间字符串，格式如 "09:30:00"
            
        返回:
            str: 时间分类（早盘板/午盘板/尾盘板）
        """
        try:
            if pd.isna(time_str) or time_str == '':
                return "未知"
            
            # 提取小时和分钟
            time_parts = str(time_str).split(':')
            if len(time_parts) < 2:
                return "未知"
            
            hour = int(time_parts[0])
            minute = int(time_parts[1])
            
            # 转换为分钟数
            total_minutes = hour * 60 + minute
            
            # 分类
            if total_minutes < 600:  # 10:00之前
                return "早盘板"
            elif total_minutes < 810:  # 13:30之前
                return "午盘板"
            else:
                return "尾盘板"
                
        except Exception as e:
            logger.error(f"解析涨停时间失败: {e}")
            return "未知"
    
    @staticmethod
    def analyze_limit_up_stock(code: str, df: pd.DataFrame) -> Dict:
        """
        分析单只涨停股票
        
        参数:
            code: 股票代码
            df: 涨停板数据DataFrame
            
        返回:
            Dict: 分析结果
        """
        try:
            stock_data = df[df['code'] == code].iloc[0]
            
            # 计算封板强度
            seal_strength = LimitUpDataFetcher.calculate_seal_strength(stock_data)
            
            # 涨停时间分类
            time_category = LimitUpDataFetcher.get_limit_up_time_category(
                stock_data.get('first_limit_up_time', '')
            )
            
            # 连板情况
            consecutive = int(stock_data.get('consecutive_limit_up', 0))
            
            result = {
                'code': code,
                'name': stock_data.get('name', ''),
                'price': stock_data.get('price', 0),
                'change_pct': stock_data.get('change_pct', 0),
                'seal_strength': seal_strength,
                'time_category': time_category,
                'consecutive': consecutive,
                'first_limit_time': stock_data.get('first_limit_up_time', ''),
                'last_limit_time': stock_data.get('last_limit_up_time', ''),
                'turnover_rate': stock_data.get('turnover_rate', 0),
                'amount': stock_data.get('amount', 0),
                'industry': stock_data.get('industry', ''),
                'seal_amount': stock_data.get('seal_amount', 0),
            }
            
            # 评分
            score = 0
            if seal_strength >= 5:
                score += 30
            elif seal_strength >= 2:
                score += 20
            else:
                score += 10
            
            if time_category == "早盘板":
                score += 30
            elif time_category == "午盘板":
                score += 20
            else:
                score += 10
            
            if consecutive >= 3:
                score += 40
            elif consecutive == 2:
                score += 30
            elif consecutive == 1:
                score += 20
            else:
                score += 10
            
            result['score'] = min(score, 100)
            
            return result
            
        except Exception as e:
            logger.error(f"分析涨停股票{code}失败: {e}")
            return {}


class LimitUpHistory:
    """涨停板历史数据分析"""
    
    @staticmethod
    def get_next_day_performance(codes: List[str], limit_up_date: str) -> pd.DataFrame:
        """
        获取涨停股票次日表现
        
        参数:
            codes: 股票代码列表
            limit_up_date: 涨停日期
            
        返回:
            DataFrame: 次日表现数据
        """
        results = []
        
        for code in codes:
            try:
                # 获取历史数据
                df = ak.stock_zh_a_hist(
                    symbol=code,
                    period="daily",
                    start_date=limit_up_date,
                    adjust="qfq"
                )
                
                if df is not None and len(df) >= 2:
                    # 涨停日和次日数据
                    limit_day = df.iloc[0]
                    next_day = df.iloc[1]
                    
                    # 计算次日涨跌幅
                    next_day_change = ((next_day['收盘'] - limit_day['收盘']) / limit_day['收盘']) * 100
                    
                    results.append({
                        'code': code,
                        'limit_up_date': limit_up_date,
                        'limit_up_price': limit_day['收盘'],
                        'next_day_change': next_day_change,
                        'next_day_high': ((next_day['最高'] - limit_day['收盘']) / limit_day['收盘']) * 100,
                        'next_day_low': ((next_day['最低'] - limit_day['收盘']) / limit_day['收盘']) * 100,
                    })
                    
            except Exception as e:
                logger.error(f"获取{code}次日表现失败: {e}")
                continue
        
        return pd.DataFrame(results)

